from typing import TypedDict
from langgraph.graph import StateGraph
from .nodes import NovelNodes

class NovelState(TypedDict):
    outline: str
    iteration: int
    current_chapter: str

class NovelWorkflow:
    def __init__(self, nodes: NovelNodes):
        self.nodes = nodes
        self.graph = StateGraph(NovelState)
    
    def _should_refine_outline(self, state):
        return state["iteration"] < 5
    
    def compile_workflow(self):
        # 添加核心节点
        self.graph.add_node("generate_outline", self.nodes.generate_initial_outline)
        self.graph.add_node("refine_outline", self.nodes.refine_outline)
        self.graph.add_node("write_chapter", self.nodes.write_chapter)
        
        # 设置初始流程
        self.graph.set_entry_point("generate_outline")
        
        # 大纲迭代条件分支
        self.graph.add_conditional_edges(
            "generate_outline",
            self._should_refine_outline,
            {
                True: "refine_outline",
                False: "write_chapter"
            }
        )
        
        self.graph.add_edge("refine_outline", "write_chapter")
        return self.graph.compile(checkpointer=None, interrupt_before=None, interrupt_after=None, debug=False, recursion_limit=5)